import random


class Node(object):
    def __init__(self, val):
        self.val = val
        self.left: Node = None
        self.right: Node = None
        self.height = 1


class AVLTree(object):
    def __init__(self):
        self.root = None

    def insert(self, val):
        self.root = self._insert(self.root, val)

    # 辅助插入函数
    def _insert(self, node: Node, val) -> Node:
        # 二叉平衡树插入
        if not node:
            return Node(val)

        if val < node.val:
            node.left = self._insert(node.left, val)

        else:
            node.right = self._insert(node.right, val)
        # 左右子树可能有变化，及时更新
        node.height = self._update_height(node)

        # 左加右，知道那边更高
        balance = self._get_balance(node)

        # 情况一
        if balance > 1 and val < node.left.val:
            return self._rotate_right(node)

        # 情况二
        if balance < -1 and val >= node.right.val:
            return self._rotate_left(node)

        # 情况三
        if balance > 1 and val >= node.left.val:
            node.left = self._rotate_left(node.left)
            return self._rotate_right(node)

        # 情况四
        if balance < -1 and val < node.right.val:
            node.right = self._rotate_right(node.right)
            return self._rotate_left(node)

        return node

    def delete(self, val):
        self.root = self._delete(self.root, val)

    def _delete(self, node: Node, val) -> Node:
        if not node:
            return node

        if val < node.val:
            node.left = self._delete(node.left, val)
        elif val > node.val:
            node.right = self._delete(node.right, val)
        else:
            # 情况一和二
            if not node.left or not node.right:
                temp = node.left if node.left else node.right
                del node
                # 情况一
                if not temp:
                    return None
                # 情况二
                else:
                    return temp
            # 情况三
            else:
                # 找到后驱节点值
                right_val = self._find_min_left_val(node.right)
                # 替换
                node.val = right_val
                # 删除后驱节点
                node.right = self._delete(node.right, right_val)

        node.height = self._update_height(node)

        balance = self._get_balance(node)

        # 判断是LL型
        if balance > 1 and self._get_balance(node.left) >= 0:
            return self._rotate_right(node)
        # RR型
        if balance < -1 and self._get_balance(node.right) <= 0:
            return self._rotate_left(node)
        # LR型
        if balance > 1 and self._get_balance(node.left) < 0:
            node.left = self._rotate_left(node.left)
            return self._rotate_right(node)

        if balance < -1 and self._get_balance(node.right) > 0:
            node.right = self._rotate_right(node.right)
            return self._rotate_left(node)

        return node

    def _find_min_left_val(self, node: Node) -> int:
        while node.left:
            node = node.left
        return node.val

    def _get_height(self, node: Node) -> int:
        if not node:
            return 0
        return node.height

    def _get_balance(self, node: Node) -> int:
        if not node:
            return 0
        return self._get_height(node.left) - self._get_height(node.right)

    def _update_height(self, node: Node) -> int:

        return 1 + max(self._get_height(node.left), self._get_height(node.right))

    def _rotate_right(self, node: Node) -> Node:
        l = node.left
        l_r = l.right

        node.left = l_r
        l.right = node

        # 记得及时更新，更新顺序不能调换！！！！
        node.height = self._update_height(node)
        l.height = self._update_height(l)

        return l
        pass

    def _rotate_left(self, node: Node) -> Node:
        r = node.right
        r_l = r.left

        node.right = r_l
        r.left = node

        node.height = self._update_height(node)
        r.height = self._update_height(r)

        return r

        pass

    def __str__(self):
        # 打印函数
        return str(self.inorder())

    # 中序遍历
    def inorder(self):
        inorder = []
        self._inorder(self.root, inorder)
        return inorder

    def _inorder(self, node: Node, inorder: list):
        if not node:
            return
        self._inorder(node.left, inorder)
        inorder.append(node.val)
        self._inorder(node.right, inorder)


def test():
    # 判断是否平衡
    def judge_balance(node: Node) -> int:
        if not node:
            return 0
        left = judge_balance(node.left)
        right = judge_balance(node.right)
        if left == -1 or right == -1:
            return -1
        diff = abs(left - right)
        if diff > 1:
            return -1
        return 1 + max(left, right)

    # 判断是否从小到大排序
    def judge_sort(arr: list) -> bool:
        for i in range(1, len(arr)):
            if arr[i] < arr[i - 1]:
                return False
        return True
    # 判断被删除节点是否正确被删除
    def judge_is_delete(arr, val):
        l, r = 0, len(arr) - 1
        while l < r:
            mid = (l + r) // 2
            if arr[mid] == val:
                return False
            elif arr[mid] < val:
                l = mid + 1
            else:
                r = mid
        return True

    # 大量数据测试
    for len_ in [random.randint(1, 2000) for _ in range(100)]:
        avl_tree = AVLTree()
        # 将数组打乱
        vals = random.sample(list(range(len_)), len_)
        for i in vals:
            avl_tree.insert(i)
        # 插入是否正常
        if not judge_sort(avl_tree.inorder()) or judge_balance(avl_tree.root) == -1:
            print(vals)
            print('error')
            break

        # 删除是否正常
        val = random.randint(0, len_ - 1)
        avl_tree.delete(val)

        if judge_balance(avl_tree.root) == -1 or not judge_is_delete(avl_tree.inorder(), val):
            print(vals, val)
            print('delete error')


if __name__ == '__main__':
    # avl_tree = AVLTree()
    # vals = random.sample(list(range(4)), 4)
    # print(vals)
    # for i in vals:
    #     avl_tree.insert(i)
    # print(avl_tree)
    # avl_tree.delete(random.randint(0, 3))
    #
    # print(avl_tree)

    test()
